import os
import json
import logging
import subprocess
import random
import socket
from typing import Dict, List, Any, Optional
from jinja2 import Environment, FileSystemLoader, select_autoescape
from flask import current_app

logger = logging.getLogger("mobox-dashboard.docker-experiment")

class DockerExperimentService:
    """Docker实验服务类，用于管理实验相关的Docker容器
    
    从参考项目 mobox_server 学习并集成的功能：
    1. 动态端口分配
    2. Docker Compose 模板管理
    3. 实验环境启动/停止
    """
    
    def __init__(self, app=None):
        """初始化Docker实验服务
        
        Args:
            app: Flask应用实例，可为None
        """
        self.app = app
        self.active_ports = set()
        
        if app is not None:
            self.init_app(app)
    
    def init_app(self, app):
        """初始化应用配置
        
        Args:
            app: Flask应用实例
        """
        self.app = app
        
        # 设置Jinja2模板环境
        template_dir = os.path.join(app.instance_path, "templates")
        if not os.path.exists(template_dir):
            template_dir = os.path.join(os.path.dirname(__file__), "..", "..", "templates")
        
        self.jinja_env = Environment(
            loader=FileSystemLoader(template_dir),
            autoescape=select_autoescape(['yml', 'yaml'])
        )
        
        self._load_active_ports()
    
    def _load_active_ports(self):
        """加载当前正在使用的端口
        
        通过检查运行中的 Docker 容器来获取已占用的端口
        """
        try:
            # 使用docker ps获取当前运行的容器端口信息
            result = subprocess.run(
                ["docker", "ps", "--format", "{{.Ports}}"],
                capture_output=True,
                text=True,
                check=True
            )
            
            ports = set()
            for line in result.stdout.splitlines():
                if "->" in line:
                    # 解析端口映射，格式类似 0.0.0.0:8080->80/tcp
                    port_mappings = line.split(",")
                    for mapping in port_mappings:
                        if "->" in mapping:
                            host_port = mapping.split("->")[0].split(":")[-1]
                            try:
                                ports.add(int(host_port))
                            except ValueError:
                                continue
            
            self.active_ports = ports
            logger.info(f"已加载活跃端口: {len(self.active_ports)} 个")
            
        except subprocess.CalledProcessError as e:
            logger.error(f"获取Docker容器信息失败: {str(e)}")
            self.active_ports = set()
        except Exception as e:
            logger.error(f"加载活跃端口时出错: {str(e)}", exc_info=True)
            self.active_ports = set()
    
    def _get_free_port(self) -> int:
        """获取一个空闲的端口
        
        Returns:
            int: 可用的端口号
            
        Raises:
            ValueError: 如果没有可用端口
        """
        port_range_start = current_app.config.get('PORT_RANGE_START', 8001)
        port_range_end = current_app.config.get('PORT_RANGE_END', 9000)
        
        port_range = range(port_range_start, port_range_end + 1)
        available_ports = list(set(port_range) - self.active_ports)
        
        if not available_ports:
            raise ValueError("没有可用的端口")
        
        # 随机选择一个可用端口
        port = random.choice(available_ports)
        
        # 检查端口是否真的可用
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            if s.connect_ex(('localhost', port)) == 0:
                # 端口已被占用，从可用列表中移除并重试
                available_ports.remove(port)
                if available_ports:
                    port = random.choice(available_ports)
                else:
                    raise ValueError("没有可用的端口")
        
        self.active_ports.add(port)
        return port
    
    def generate_docker_compose(self, 
                               user_id: str, 
                               experiment_id: str, 
                               course_name: str, 
                               port: int,
                               user_dir: str,
                               course_files_path: str) -> str:
        """生成docker-compose配置文件
        
        Args:
            user_id: 用户ID
            experiment_id: 实验ID
            course_name: 课程名称
            port: 端口号
            user_dir: 用户目录
            course_files_path: 课程文件路径
            
        Returns:
            str: 生成的docker-compose文件路径
        """
        try:
            template = self.jinja_env.get_template("docker-compose.yml.j2")
        except Exception as e:
            logger.error(f"无法加载Docker Compose模板: {str(e)}")
            # 如果模板不存在，使用默认配置
            return self._generate_default_docker_compose(
                user_id, experiment_id, course_name, port, user_dir, course_files_path
            )
        
        service_name = f"{user_id}_{course_name}_{experiment_id}"
        
        # 构建模板上下文
        context = {
            "service_name": service_name,
            "user_id": user_id,
            "experiment_id": experiment_id,
            "course_name": course_name,
            "port": port,
            "user_dir": user_dir,
            "course_files_path": course_files_path,
            "host_ip": current_app.config.get('HOST_IP', 'localhost')
        }
        
        # 渲染模板
        docker_compose_content = template.render(**context)
        
        # 保存到文件
        output_dir = os.path.join(user_dir, "docker_compose")
        os.makedirs(output_dir, exist_ok=True)
        output_file = os.path.join(output_dir, f"{service_name}.yml")
        
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write(docker_compose_content)
        
        logger.info(f"已生成docker-compose文件: {output_file}")
        return output_file
    
    def _generate_default_docker_compose(self, 
                                       user_id: str, 
                                       experiment_id: str, 
                                       course_name: str, 
                                       port: int,
                                       user_dir: str,
                                       course_files_path: str) -> str:
        """生成默认的docker-compose配置
        
        Args:
            user_id: 用户ID
            experiment_id: 实验ID
            course_name: 课程名称
            port: 端口号
            user_dir: 用户目录
            course_files_path: 课程文件路径
            
        Returns:
            str: 生成的docker-compose文件路径
        """
        service_name = f"{user_id}_{course_name}_{experiment_id}"
        
        # 默认的docker-compose配置
        docker_compose_content = f"""version: '3.8'

services:
  {service_name}:
    image: jupyter/base-notebook:latest
    container_name: {service_name}
    ports:
      - "{port}:8888"
    volumes:
      - "{user_dir}:/home/jovyan/work"
      - "{course_files_path}:/home/jovyan/course"
    environment:
      - JUPYTER_ENABLE_LAB=yes
      - JUPYTER_TOKEN=""
    command: start-notebook.sh --NotebookApp.token='' --NotebookApp.password=''
    restart: unless-stopped
"""
        
        # 保存到文件
        output_dir = os.path.join(user_dir, "docker_compose")
        os.makedirs(output_dir, exist_ok=True)
        output_file = os.path.join(output_dir, f"{service_name}.yml")
        
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write(docker_compose_content)
        
        logger.info(f"已生成默认docker-compose文件: {output_file}")
        return output_file
    
    def start_experiment_service(self, 
                                user_id: str, 
                                experiment_id: str, 
                                course_name: str,
                                user_dir: str,
                                course_files_path: str) -> Dict[str, Any]:
        """启动实验服务
        
        Args:
            user_id: 用户ID
            experiment_id: 实验ID
            course_name: 课程名称
            user_dir: 用户目录路径
            course_files_path: 课程文件路径
            
        Returns:
            Dict[str, Any]: 包含服务信息的字典
        """
        try:
            # 分配端口
            port = self._get_free_port()
            
            # 生成docker-compose文件
            docker_compose_file = self.generate_docker_compose(
                user_id=user_id,
                experiment_id=experiment_id,
                course_name=course_name,
                port=port,
                user_dir=user_dir,
                course_files_path=course_files_path
            )
            
            # 获取服务名称
            service_name = f"{user_id}_{course_name}_{experiment_id}"
            
            # 使用docker-compose启动服务
            logger.info(f"开始启动实验服务: {service_name}")
            result = subprocess.run(
                ["docker-compose", "-f", docker_compose_file, "-p", service_name, "up", "-d"],
                capture_output=True,
                text=True,
                check=True
            )
            
            logger.info(f"实验服务启动成功: {service_name}")
            logger.debug(f"Docker输出: {result.stdout}")
            
            # 构建服务访问URL
            host_ip = current_app.config.get('HOST_IP', 'localhost')
            service_url = f"http://{host_ip}:{port}/"
            
            return {
                'success': True,
                'service_name': service_name,
                'port': port,
                'url': service_url,
                'docker_compose_file': docker_compose_file
            }
            
        except subprocess.CalledProcessError as e:
            error_msg = f"启动Docker服务失败: {e.stderr}"
            logger.error(error_msg, exc_info=True)
            return {
                'success': False,
                'error': error_msg
            }
        except Exception as e:
            error_msg = f"启动实验服务时发生错误: {str(e)}"
            logger.error(error_msg, exc_info=True)
            return {
                'success': False,
                'error': error_msg
            }
    
    def stop_experiment_service(self, user_id: str, experiment_id: str, course_name: str) -> Dict[str, Any]:
        """停止实验服务
        
        Args:
            user_id: 用户ID
            experiment_id: 实验ID
            course_name: 课程名称
            
        Returns:
            Dict[str, Any]: 操作结果
        """
        service_name = f"{user_id}_{course_name}_{experiment_id}"
        
        try:
            # 使用docker-compose停止服务
            result = subprocess.run(
                ["docker-compose", "-p", service_name, "down"],
                capture_output=True,
                text=True,
                check=True
            )
            
            logger.info(f"实验服务停止成功: {service_name}")
            return {
                'success': True,
                'service_name': service_name,
                'message': '服务已成功停止'
            }
            
        except subprocess.CalledProcessError as e:
            error_msg = f"停止Docker服务失败: {e.stderr}"
            logger.error(error_msg)
            return {
                'success': False,
                'error': error_msg
            }
        except Exception as e:
            error_msg = f"停止实验服务时发生错误: {str(e)}"
            logger.error(error_msg, exc_info=True)
            return {
                'success': False,
                'error': error_msg
            }
    
    def list_user_experiments(self, user_id: str) -> List[Dict[str, Any]]:
        """列出用户的实验服务
        
        Args:
            user_id: 用户ID
            
        Returns:
            List[Dict[str, Any]]: 用户实验服务列表
        """
        try:
            # 获取所有Docker容器
            result = subprocess.run(
                ["docker", "ps", "-a", "--format", "json"],
                capture_output=True,
                text=True,
                check=True
            )
            
            experiments = []
            for line in result.stdout.strip().split('\n'):
                if line:
                    container_info = json.loads(line)
                    container_name = container_info.get('Names', '')
                    
                    # 检查是否是该用户的实验容器
                    if container_name.startswith(f"{user_id}_"):
                        parts = container_name.split('_')
                        if len(parts) >= 3:
                            course_name = '_'.join(parts[1:-1])
                            experiment_id = parts[-1]
                            
                            experiments.append({
                                'experiment_id': experiment_id,
                                'course_name': course_name,
                                'status': container_info.get('State', 'unknown'),
                                'ports': container_info.get('Ports', ''),
                                'created': container_info.get('CreatedAt', ''),
                                'container_name': container_name
                            })
            
            return experiments
            
        except subprocess.CalledProcessError as e:
            logger.error(f"获取实验列表失败: {str(e)}")
            return []
        except Exception as e:
            logger.error(f"列出用户实验时发生错误: {str(e)}", exc_info=True)
            return [] 